%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Program: marginal_utilities.m
% By: Stephie Fried, Kevin Novan, and William Peterman
% Date: Fall 2024
% Purpose: Calculates marginal utility as a function of total income and labor
% income 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
clc;
close all;
clear all;

bpath = 'C:\Users\l1sxf17\Dropbox (FRB SF)\Research\Will\equity_efficiency_project\Replication';
cd(bpath);

options = optimset('Algorithm','Levenberg-Marquardt','TolFun',10e-15,'TolX',10e-15, ...
    'MaxFunEvals', 10000, 'MaxIter', 10000, 'Display','none');

%% load data
% Parameters 
cd(bpath);
load('Programs/Matlab/aggregates')

%reshape pop
people = zeros(810000,3); %[person, age, weight]
age = zeros(810000,1);
person =age; 
j=1;
k=1;
for i =1:810000
    people(i,1)=k;
    people(i,2)=j;
    if j==81 %J = 81, start over, go to next person
        j=1;
        k=k+1;
    else
        j=j+1;
    end
    age(i) =j-1;
    person(i) =k;
    people(i,3)=popweights(people(i,2));
end
popweights2 = people(:,3)/10000; 

%Marginal utility in the baseline 
c_raw =importdata('Fortran_output/baseline/sim_consumption_output.dat');
ec_raw =importdata('Fortran_output/baseline/sim_energy_output.dat');
capital_raw =importdata('Fortran_output/baseline/sim_capital_output.dat');
earnings_raw =importdata('Fortran_output/baseline/sim_labor_earnings_output.dat');
labor_taxes_raw = importdata('Fortran_output/baseline/sim_labor_tax_output.dat');
taxable_earnings_raw = earnings_raw - 0.5*tauss_base*min(earnings_raw, 2.42*avg_earnings_base);

%Reshape capital to get rid of 82nd year in capital
counter =0; 
counter2=0;
capital_raw2 = zeros(size(c_raw));
for i = 1:10000
    for j = 1:J+1
        counter = counter+1;
        if j < J+1
            counter2 = counter2+1;
            capital_raw2(counter2) = capital_raw(counter);
        end
    end
end

ctilde_raw = c_raw.^gamma.*(ec_raw - ebar).^(1-gamma); 
mu_raw = theta1*ctilde_raw.^(-theta1);
income_raw = earnings_raw + r*capital_raw2; 
taxable_income_raw = taxable_earnings_raw + r*capital_raw2; 
capital_taxes_raw = capital_raw2*r*lambdak;
capital_taxes_raw(capital_taxes_raw <0) =0; 
capital_income_raw = r*capital_raw2;

%Optimal 
labor_taxes_opt_raw = importdata('Fortran_output/optimal/sim_labor_tax_new_optimal.dat');
labor_taxes_opt_old_raw = importdata('Fortran_output/optimal/sim_labor_tax_old_optimal.dat');
capital_taxes_opt_raw = importdata('Fortran_output/optimal/sim_capital_output_optimal.dat')*OPT(r_ind)*OPT(lambdak_ind);
capital_taxes_opt_raw(capital_taxes_opt_raw <0) =0;

%Reshape capital taxes to get rid of 82nd year in capital
counter =0; 
counter2=0;
capital_taxes_opt_raw2 = zeros(size(c_raw));
for i = 1:10000
    for j = 1:J+1
        counter = counter+1;
        if j < J+1
            counter2 = counter2+1;
            capital_taxes_opt_raw2(counter2) = capital_taxes_opt_raw(counter);
        end
    end
end

%Income dependent rebate (maximize equity)
inc_dep_rebate_me_raw = importdata('Fortran_output/income_dependent/sim_rebate_output_kevin_max.dat');
total_inc_dep_rebate_me = inc_dep_rebate_me_raw'*popweights2;

%Simple progressivity
labor_taxes_prog_raw = importdata('Fortran_output/prog/sim_labor_tax_new_optimal.dat');
labor_taxes_prog_old_raw = importdata('Fortran_output/prog/sim_labor_tax_old_optimal.dat');
simple_prog_rebate_raw = labor_taxes_prog_old_raw - labor_taxes_prog_raw;
simple_prog_rebate_raw(simple_prog_rebate_raw <0)=0;
total_simple_prog_rebate = simple_prog_rebate_raw'*popweights2;

%% Reformat for figure

total_lump_rebate = 2; 
%Sort people by marginal utility
mat = [popweights2, mu_raw, income_raw, earnings_raw, capital_raw2*r, inc_dep_rebate_me_raw, simple_prog_rebate_raw, ...
    ones(size(simple_prog_rebate_raw))*total_lump_rebate];
mat_sorted = sortrows(mat, 2, 'descend');

%Cumulative pop and rebate shares
cumpop = zeros(810000,1);
cs_idme = zeros(810000,1);
cs_idme(1) = mat_sorted(1, 6)*mat_sorted(1,1)/total_inc_dep_rebate_me;
cs_prog = zeros(810000,1);
cs_prog(1) = mat_sorted(1, 7)*mat_sorted(1,1)/total_simple_prog_rebate;
cs_lump = zeros(810000,1);
cs_lump(1) = mat_sorted(1, 8)*mat_sorted(1,1)/total_lump_rebate;

cumpop(1) = mat_sorted(1,1);
for i = 2:1:810000
    cumpop(i) = mat_sorted(i, 1) + cumpop(i-1);
    cs_idme(i) = cs_idme(i-1) + mat_sorted(i, 6)*mat_sorted(i,1)/total_inc_dep_rebate_me;
    cs_prog(i) = cs_prog(i-1) + mat_sorted(i, 7)*mat_sorted(i,1)/total_simple_prog_rebate;
    cs_lump(i) = cs_lump(i-1) +mat_sorted(i,8)*mat_sorted(i,1)/total_lump_rebate; 
end
save ('Programs/Matlab/mu');

